from pymilvus import Collection, FieldSchema, CollectionSchema, DataType, utility

COLLECTION_NAME = "doc_vectors"


def init_doc_collection():
    if not utility.has_collection(COLLECTION_NAME):
        fieldSchemas = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=2560),
            FieldSchema(name="doc", dtype=DataType.VARCHAR, max_length=4096),
        ]
        schema = CollectionSchema(fieldSchemas)
        # 创建集合对象
        collection = Collection(COLLECTION_NAME, schema)
        index_params = {
            "index_type": "IVF_FLAT",
            "metric_type": "COSINE",
            "params": {"nlist": 128},
        }
        collection.create_index(field_name="embedding", index_params=index_params)
        collection.load()
    else:
        collection = Collection(COLLECTION_NAME)
        collection.load()


# 插入文档向量和原文件的内容到集合去
def insert_doc_vectors(embeddings, docs):
    collection = Collection(COLLECTION_NAME)
    collection.insert([embeddings, docs])


def insert_doc_vector(embedding, doc):
    collection = Collection(COLLECTION_NAME)
    collection.insert({"embedding": embedding, "doc": doc})


def search_doc_vectors(query_embedding, limit=5):
    collection = Collection(COLLECTION_NAME)
    results = collection.search(
        data=[query_embedding],
        anns_field="embedding",
        param={"metric_type": "COSINE", "params": {"nprobe": 16}},
        limit=limit,
        output_fields=["doc"],
    )
    return results
